package dfxnftables

import (
	"context"
	"fmt"
	"net"

	"github.com/google/nftables"
	"github.com/google/nftables/expr"
)

type NfChainRule struct {
	NfMeta
	ChainName string
	// 匹配条件
	SrcSet   string   // 源IP集合名（在 TableName 下）
	DstSet   string   // 目的IP集合名（可选）
	SrcIP    string   // 单个源IP（与 SrcSet 二选一）
	DstIP    string   // 单个目的IP
	Protocol string   // tcp/udp/icmp/icmpv6/all(空视为all)
	DstPorts []uint16 // 目标端口（仅在 tcp/udp 时有效）
	// 动作与标注
	Verdict string // accept/drop/reject/return  accept允许访问， drop拒绝访问， reject拒绝访问并返回拒绝信息， return返回到上级链
	Comment string // 规则备注（用于删除查询）
}

// ensureComment 生成或返回已有的 Comment，保证规则写入时带有可追踪的标识。
func (n *NfChainRule) ensureComment() string {
	if n.Comment != "" {
		return n.Comment
	}
	comment := fmt.Sprintf("auto:%s:%s:%s:%s:%s:%s:%s:%s:%v:%s",
		n.AddressFamily,
		n.TableName,
		n.ChainName,
		n.Protocol,
		n.SrcSet,
		n.SrcIP,
		n.DstSet,
		n.DstIP,
		n.DstPorts,
		n.Verdict,
	)
	n.Comment = comment
	return comment
}

// resolveSetForLookup 根据集合名称在指定表中查找并返回对应的 nftables 集合。
// 如果集合不存在或获取失败，将返回错误。
//
// 参数:
//   - c: nftables 连接句柄
//   - tbl: 表对象
//   - setName: 要查找的集合名称
//
// 返回值:
//   - *nftables.Set: 查找到的集合对象
//   - error: 错误信息（如果有的话）
func resolveSetForLookup(c *nftables.Conn, tbl *nftables.Table, setName string) (*nftables.Set, error) {
	s, err := c.GetSetByName(tbl, setName)
	if err != nil {
		return nil, fmt.Errorf("get set %s failed: %w", setName, err)
	}
	if s == nil {
		return nil, fmt.Errorf("set %s not found in table %s", setName, tbl.Name)
	}
	if s.Name == "" {
		s.Name = setName
	}
	return s, nil
}

// AddRule 向指定的 nftables 表和链中添加一条规则。
// 支持基于协议、源IP/集合、目的IP/集合、目标端口等条件进行匹配，并设置动作（accept/drop）。
//
// 参数:
//   - ctx: 上下文
//   - c: nftables 连接句柄
//   - t: 表对象（可为 nil，此时会自动解析）
//
// 返回值:
//   - error: 添加规则过程中发生的错误（如果有）
func (n *NfChainRule) AddRule(ctx context.Context, c *nftables.Conn, tbl *nftables.Table) error {
	if err := n.NfMeta.Validate(); err != nil {
		return err
	}
	var err error
	if tbl == nil {
		tbl, err = n.NfMeta.ResolveTable(c)
		if err != nil {
			return err
		}
	}
	if tbl.Name != n.TableName {
		return fmt.Errorf("table mismatch: expect %s got %s", n.TableName, tbl.Name)
	}
	fam := n.Family()
	ch := &nftables.Chain{Name: n.ChainName, Table: tbl}

	var ex []expr.Any

	// 协议匹配（IPv4: offset 9；IPv6: offset 6）
	if n.Protocol != "" && n.Protocol != "all" {
		var protoByte uint8
		switch n.Protocol {
		case "tcp":
			protoByte = 6
		case "udp":
			protoByte = 17
		case "icmp":
			protoByte = 1
		case "icmpv6":
			protoByte = 58
		default:
			return fmt.Errorf("unsupported protocol: %s", n.Protocol)
		}
		offset := uint32(9)
		if fam == nftables.TableFamilyIPv6 {
			offset = 6
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: offset, Len: 1},
			&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{protoByte}},
		)
	}

	// 源IP集合匹配（优先于单个IP）
	if n.SrcSet != "" {
		s, err := resolveSetForLookup(c, tbl, n.SrcSet)
		if err != nil {
			return err
		}
		// 载入源地址到寄存器1：IPv4偏移12，IPv6偏移8（源地址起始）
		var off uint32 = 12
		var ln uint32 = 4
		if fam == nftables.TableFamilyIPv6 {
			off = 8
			ln = 16
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: off, Len: ln},
			&expr.Lookup{SourceRegister: 1, SetID: s.ID, SetName: s.Name},
		)
	} else if n.SrcIP != "" { // 单个源IP匹配
		ip := net.ParseIP(n.SrcIP)
		if ip == nil {
			return fmt.Errorf("invalid src ip: %s", n.SrcIP)
		}
		var off uint32 = 12
		var data []byte
		if fam == nftables.TableFamilyIPv6 {
			off = 8
			data = ip.To16()
			if data == nil {
				return fmt.Errorf("src ip is not ipv6: %s", n.SrcIP)
			}
		} else {
			data = ip.To4()
			if data == nil {
				return fmt.Errorf("src ip is not ipv4: %s", n.SrcIP)
			}
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: off, Len: uint32(len(data))},
			&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: data},
		)
	}

	// 目的IP集合匹配（可选）
	if n.DstSet != "" {
		s, err := resolveSetForLookup(c, tbl, n.DstSet)
		if err != nil {
			return err
		}
		var off uint32 = 16
		var ln uint32 = 4
		if fam == nftables.TableFamilyIPv6 {
			off = 24
			ln = 16
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: off, Len: ln},
			&expr.Lookup{SourceRegister: 1, SetID: s.ID, SetName: s.Name},
		)
	} else if n.DstIP != "" { // 单个目的IP匹配
		ip := net.ParseIP(n.DstIP)
		if ip == nil {
			return fmt.Errorf("invalid dst ip: %s", n.DstIP)
		}
		var off uint32 = 16
		var data []byte
		if fam == nftables.TableFamilyIPv6 {
			off = 24
			data = ip.To16()
			if data == nil {
				return fmt.Errorf("dst ip is not ipv6: %s", n.DstIP)
			}
		} else {
			data = ip.To4()
			if data == nil {
				return fmt.Errorf("dst ip is not ipv4: %s", n.DstIP)
			}
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: off, Len: uint32(len(data))},
			&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: data},
		)
	}

	// 目的端口匹配（tcp/udp 有效）
	if len(n.DstPorts) > 0 && (n.Protocol == "tcp" || n.Protocol == "udp") {
		// L4 目的端口偏移2，长度2
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseTransportHeader, Offset: 2, Len: 2},
		)
		if len(n.DstPorts) == 1 {
			p := n.DstPorts[0]
			ex = append(ex, &expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{byte(p >> 8), byte(p & 0xff)}})
		} else {
			// 多端口：简单起见，生成多个规则（调用方可循环调用），这里先只支持单端口
			return fmt.Errorf("multiple dst ports not supported in one rule; create multiple rules")
		}
	}

	// 动作
	switch n.Verdict {
	case "accept":
		ex = append(ex, &expr.Verdict{Kind: expr.VerdictAccept})
	case "drop":
		ex = append(ex, &expr.Verdict{Kind: expr.VerdictDrop})
	default:
		return fmt.Errorf("unsupported verdict: %s", n.Verdict)
	}

	r := &nftables.Rule{Table: tbl, Chain: ch, Exprs: ex}
	comment := n.ensureComment()
	if comment != "" {
		r.UserData = []byte(comment)
	}
	c.AddRule(r)
	return c.Flush()
}

// DelRule 删除指定表和链中的规则。
// 当前仅支持通过 Comment 字段来识别并删除规则。
//
// 参数:
//   - ctx: 上下文
//   - c: nftables 连接句柄
//
// 返回值:
//   - error: 删除规则过程中发生的错误（如果有）
func (n *NfChainRule) DelRule(ctx context.Context, c *nftables.Conn, tbl *nftables.Table) error {
	if err := n.NfMeta.Validate(); err != nil {
		return err
	}

	ch := &nftables.Chain{Name: n.ChainName, Table: tbl}
	rules, err := c.GetRules(tbl, ch)
	if err != nil {
		return err
	}
	// 简化：如果提供了 Comment，则按 Comment 删除匹配的规则
	comment := n.ensureComment()
	if comment != "" {
		for _, r := range rules {
			if string(r.UserData) == comment {
				c.DelRule(r)
			}
		}
		return c.Flush()
	}
	// 否则暂不支持更复杂的匹配删除
	return fmt.Errorf("DelRule requires Comment to identify the rule to delete (table=%s chain=%s)", n.TableName, n.ChainName)
}

// EasyAddRule 是 AddRule 的简化版本，主要用于调用方已持有表句柄的情况。
// 若未提供表句柄，则会尝试解析获取。
//
// 参数:
//   - ctx: 上下文
//   - c: nftables 连接句柄
//   - tbl: 表对象（可为 nil）
//
// 返回值:
//   - error: 添加规则过程中发生的错误（如果有）
func (n *NfChainRule) EasyAddRule(ctx context.Context, c *nftables.Conn, tbl *nftables.Table) error {
	if err := n.NfMeta.Validate(); err != nil {
		return err
	}
	fam := n.Family()
	var err error
	if tbl == nil {
		tbl, err = n.NfMeta.ResolveTable(c)
		if err != nil {
			return err
		}
	}
	if tbl == nil {
		return fmt.Errorf("table not provided and resolve failed")
	}
	if tbl.Family != fam || tbl.Name != n.TableName {
		return fmt.Errorf("table mismatch: expect %s/%d got %s/%d", n.TableName, fam, tbl.Name, tbl.Family)
	}
	ch := &nftables.Chain{Name: n.ChainName, Table: tbl}

	var ex []expr.Any

	// 协议匹配（IPv4: offset 9；IPv6: offset 6）
	if n.Protocol != "" && n.Protocol != "all" {
		var protoByte uint8
		switch n.Protocol {
		case "tcp":
			protoByte = 6
		case "udp":
			protoByte = 17
		case "icmp":
			protoByte = 1
		case "icmpv6":
			protoByte = 58
		default:
			return fmt.Errorf("unsupported protocol: %s", n.Protocol)
		}
		offset := uint32(9)
		if fam == nftables.TableFamilyIPv6 {
			offset = 6
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: offset, Len: 1},
			&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{protoByte}},
		)
	}

	// 源IP集合匹配（优先于单个IP）
	if n.SrcSet != "" {
		s, err := resolveSetForLookup(c, tbl, n.SrcSet)
		if err != nil {
			return err
		}
		var off uint32 = 12
		var ln uint32 = 4
		if fam == nftables.TableFamilyIPv6 {
			off = 8
			ln = 16
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: off, Len: ln},
			&expr.Lookup{SourceRegister: 1, SetID: s.ID, SetName: s.Name},
		)
	} else if n.SrcIP != "" {
		ip := net.ParseIP(n.SrcIP)
		if ip == nil {
			return fmt.Errorf("invalid src ip: %s", n.SrcIP)
		}
		var off uint32 = 12
		var data []byte
		if fam == nftables.TableFamilyIPv6 {
			off = 8
			data = ip.To16()
			if data == nil {
				return fmt.Errorf("src ip is not ipv6: %s", n.SrcIP)
			}
		} else {
			data = ip.To4()
			if data == nil {
				return fmt.Errorf("src ip is not ipv4: %s", n.SrcIP)
			}
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: off, Len: uint32(len(data))},
			&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: data},
		)
	}

	if n.DstSet != "" {
		s, err := resolveSetForLookup(c, tbl, n.DstSet)
		if err != nil {
			return err
		}
		var off uint32 = 16
		var ln uint32 = 4
		if fam == nftables.TableFamilyIPv6 {
			off = 24
			ln = 16
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: off, Len: ln},
			&expr.Lookup{SourceRegister: 1, SetID: s.ID, SetName: s.Name},
		)
	} else if n.DstIP != "" {
		ip := net.ParseIP(n.DstIP)
		if ip == nil {
			return fmt.Errorf("invalid dst ip: %s", n.DstIP)
		}
		var off uint32 = 16
		var data []byte
		if fam == nftables.TableFamilyIPv6 {
			off = 24
			data = ip.To16()
			if data == nil {
				return fmt.Errorf("dst ip is not ipv6: %s", n.DstIP)
			}
		} else {
			data = ip.To4()
			if data == nil {
				return fmt.Errorf("dst ip is not ipv4: %s", n.DstIP)
			}
		}
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseNetworkHeader, Offset: off, Len: uint32(len(data))},
			&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: data},
		)
	}

	if len(n.DstPorts) > 0 && (n.Protocol == "tcp" || n.Protocol == "udp") {
		ex = append(ex,
			&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseTransportHeader, Offset: 2, Len: 2},
		)
		if len(n.DstPorts) == 1 {
			p := n.DstPorts[0]
			ex = append(ex, &expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{byte(p >> 8), byte(p & 0xff)}})
		} else {
			return fmt.Errorf("multiple dst ports not supported in one rule; create multiple rules")
		}
	}

	switch n.Verdict {
	case "accept":
		ex = append(ex, &expr.Verdict{Kind: expr.VerdictAccept})
	case "drop":
		ex = append(ex, &expr.Verdict{Kind: expr.VerdictDrop})
	default:
		return fmt.Errorf("unsupported verdict: %s", n.Verdict)
	}

	r := &nftables.Rule{Table: tbl, Chain: ch, Exprs: ex}
	comment := n.ensureComment()
	if comment != "" {
		r.UserData = []byte(comment)
	}
	c.AddRule(r)
	return c.Flush()
}

func (r *NfChainRule) CheckRule(ctx context.Context, nftc *nftables.Conn, nt *nftables.Table, nc *nftables.Chain) (bool, error) {
	rules, err := nftc.GetRules(nt, nc)
	if err != nil {
		return false, err
	}
	comment := r.ensureComment()
	for _, rule := range rules {
		if rule.UserData != nil && string(rule.UserData) == comment {
			return true, nil
		}

	}

	return false, nil
}
